import matplotlib.pyplot as plt
import numpy as np

np.random.seed(42)


# 随意定义了一个2次多项式
def f(x):
    return -3 * x ** 2 + 2 * x + 20


# 设计一个返回正态随机分布的函数，用于产生噪声
# 给 x 定义的 ndarray 对象的结构，给每个对应位置生成一个随机噪声值
def noise(x, scale):
    return np.random.normal(scale=scale, size=x.shape)


# 将每个位置生成的噪声值添加到 x 这个 ndarray 对象中的每个值上。这是应用噪声改变数值
def add_noise(x, mult, add):
    return x * (1 + noise(x, mult)) + noise(x, add)


# 生成指定范围内的20个随机数
x = np.linspace(-2, 2, num=20)[:, None]
# 将随机数作为参数传入多项式函数计算结果，并对结果生成随机噪声
y = add_noise(f(x), 0.2, 1.3)

# 绘制散点图，注意是plt.show()，而不是scatter()结果对象的show()
# plt.scatter(x, y)
# plt.show()

# 根据上面生成的带噪声的的数据，现在去反向拟合1，2，3阶多项式
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import make_pipeline


# 定义一个绘图程序，绘制给定函数 f 的 x 在（-2.1, 2.1）之间的函数图像
def plot_function(f, min=-2.1, max=2.1, color='r'):
    x = np.linspace(min, max, 100)[:, None]
    plt.plot(x, f(x), color)


def plot_polynomial(degree, color='r'):
    model = make_pipeline(PolynomialFeatures(degree), LinearRegression())
    # 开始拟合
    model.fit(x, y)
    # 拟合得到的预测模型函数
    predict = model.predict
    print(predict)

    # 绘制拟合函数的图像
    plot_function(model.predict, color=color)


# 绘制用于预测的数据点位置
plt.scatter(x, y)
# 绘制原本多项式的函数图像
plot_function(f, color='r')
# 绘制一阶多项式拟合得到的函数图像，结果是 under-fit
plot_polynomial(1, color='g')
# 绘制十阶多项式拟合得到的函数图像，结果是 over-fit
plot_polynomial(2, color='b')

plt.show()
